package handlers

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"sdapiproject/models"
	"sdapiproject/redis_func"
	"sdapiproject/services"
	"sdapiproject/task_process"
	"sdapiproject/utils"
	"strconv"
	"strings"
)

// @Summary		测试SayHello
// @Description	向你说Hello
// @Tags			测试
// @Accept			mpfd
// @Produce		json
// @Param			who	query		string	true			"人名"
// @Success		200	{string}	string	"{"msg": "hello	Razeen"}"
// @Failure		400	{string}	string	"{"msg": "who	are	you"}"
// @Router			/hello [get]
func HandleHello(c *gin.Context) {
	who := c.Query("who")

	if who == "" {
		c.JSON(http.StatusBadRequest, gin.H{"msg": "who are u?"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"msg": "hello " + who})
}

// @Summary		文本转图片请求
// @Description	文本转图片请求
// @Tags			文本转图片
// @Accept			mpfd
// @Produce		json
// @Param			prompt	formData	string	true	"提示词"
// @Success		200		{object}	object	"{"task_id": "task_id"}"
// @Failure		400		{object}	object	"{"error": "prompt	is	empty"}"
// @Router			/sdtxt2img [post]
func Txt2ImgRequest(c *gin.Context) {
	//data := models.NewData()
	var err error

	//data.SendImages = true
	//data.TaskId = utils.GenerateTaskID()
	task_id := utils.GenerateTaskID()
	//判断propmt是否为空
	if c.PostForm("prompt") == "" {
		redis_func.SetTxt2ImgStatusToRedis(task_id, redis_func.FAILED)
		c.JSON(http.StatusBadRequest, utils.ErrorResponse(http.StatusBadRequest, "prompt is empty"))
		return
	}
	//data.Prompt = c.PostForm("prompt")
	prompt := c.PostForm("prompt")

	//data.Prompt, err = utils.TranslateToENOne(data.Prompt)

	redis_func.SetTxt2ImgStatusToRedis(task_id, redis_func.QUEUED)

	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(task_id, redis_func.FAILED)
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
		return
	}

	//url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/txt2img"

	//jsonData, err := json.Marshal(data)
	if err != nil {
		redis_func.SetTxt2ImgStatusToRedis(task_id, redis_func.FAILED)
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
		return
	}

	task_process.AddTxt2ImgTask(prompt, task_id)
	// 提前返回 task_id
	c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"task_id": task_id}))

	//// 将后续处理放入新的 goroutine 中执行
	//go func() {
	//
	//	resp, err := services.SubmitPost(url, jsonData)
	//	if err != nil {
	//		redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.FAILED)
	//		return
	//	}
	//
	//	defer resp.Body.Close()
	//
	//	body, err := ioutil.ReadAll(resp.Body)
	//	if err != nil {
	//		redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.FAILED)
	//		log.Println(err)
	//		return
	//	}
	//
	//	var response models.SuccessResponse
	//
	//	err = json.Unmarshal(body, &response)
	//	if err != nil {
	//		redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.FAILED)
	//		log.Println(err)
	//		return
	//	}
	//
	//	dir, err := utils.CreateTaskDir(data.TaskId)
	//	if err != nil || dir == "" {
	//		redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.FAILED)
	//		log.Println(err)
	//		return
	//	}
	//
	//	for i, img := range response.Images {
	//
	//		name := filepath.Join(dir, "img"+strconv.Itoa(i+1)+".png")
	//		err = utils.SaveEncodedImage(img, name)
	//		if err != nil {
	//			log.Println(err)
	//			return
	//		}
	//		response.Images[i] = name
	//	}
	//
	//	result, _ := json.Marshal(response)
	//
	//	err = redis_func.SetTxt2ImgTaskToRedis(data.TaskId, string(result))
	//	if err != nil {
	//		log.Println(err)
	//		return
	//	}
	//	redis_func.SetTxt2ImgStatusToRedis(data.TaskId, redis_func.SUCCESS)
	//
	//}()
}

func Txt2ImgRequest_all(c *gin.Context) {
	var data models.Data
	var err error

	if err = c.ShouldBindJSON(&data); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	data.SendImages = true
	data.TaskId = utils.GenerateTaskID()

	data.Prompt, err = utils.TranslateToENOne(data.Prompt)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/txt2img"

	jsonData, err := json.Marshal(data)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// 提前返回 task_id
	c.JSON(http.StatusOK, gin.H{"task_id": data.TaskId})

	// 将后续处理放入新的 goroutine 中执行
	go func() {

		resp, err := services.SubmitPost(url, jsonData)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}

		defer resp.Body.Close()

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Println(err)
			return
		}

		var response models.SuccessResponse

		err = json.Unmarshal(body, &response)
		if err != nil {
			log.Println(err)
			return
		}

		dir, err := utils.CreateTaskDir(data.TaskId)
		if err != nil || dir == "" {
			log.Println(err)
			return
		}

		for i, img := range response.Images {

			name := filepath.Join(dir, "img"+strconv.Itoa(i+1)+".png")
			err = utils.SaveEncodedImage(img, name)
			if err != nil {
				log.Println(err)
				return
			}
			response.Images[i] = name
		}

		result, _ := json.Marshal(response)

		err = redis_func.SetTxt2ImgTaskToRedis(data.TaskId, string(result))
		if err != nil {
			log.Println(err)
			return
		}
	}()
}

func Txt2ImgRequest_bak(c *gin.Context) {
	var data models.Data
	var err error

	if err = c.ShouldBindJSON(&data); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	data.SendImages = true
	data.TaskId = utils.GenerateTaskID()

	data.Prompt, err = utils.TranslateToENOne(data.Prompt)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/txt2img"

	jsonData, err := json.Marshal(data)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	resp, err := services.SubmitPost(url, jsonData)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	var response models.SuccessResponse

	err = json.Unmarshal(body, &response)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
	}

	dir, err := utils.CreateTaskDir(data.TaskId)
	if err != nil || dir == "" {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})

	}

	for i, img := range response.Images {

		name := filepath.Join(dir, "img"+strconv.Itoa(i+1)+".png")
		err = utils.SaveEncodedImage(img, name)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		response.Images[i] = name
	}

	//err = utils.SaveEncodedImage(response.Images[0], "dog.png")
	//if err != nil {
	//	c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
	//	return
	//}

	result, _ := json.Marshal(response)

	err = redis_func.SetTxt2ImgTaskToRedis(data.TaskId, string(result))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	c.JSON(http.StatusOK, gin.H{"task_id": data.TaskId})

	//c.JSON(http.StatusOK, response)

}

// @Summary		获取进度
// @Description	获取进度
// @Tags			文本转图片
// @Accept			mpfd
// @Produce		json
// @Param			task_id	formData	string	true	"任务ID"
// @Success		200		{string}	string	""
// @Failure		400		{string}	string	""
// @Router			/sdtxt2img/progress [post]
func ProgressRequest(c *gin.Context) {
	url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/progress"

	taskId := c.PostForm("task_id")
	taskStatus, err2 := redis_func.GetTxt2ImgStatusFromRedis(taskId)
	if err2 != nil {
		fmt.Println("获取进度失败", err2.Error())
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err2.Error()))
		return
	}
	if taskStatus == redis_func.SUCCESS {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task success"}))
		return
	} else if taskStatus == redis_func.FAILED {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task failed"}))
		return
	}

	url = url + "?txt2img_task_id=" + taskId
	resp, err := services.SubmitPost(url, nil)
	if err != nil {
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
	}
	var response models.ProcessResponse

	err = json.Unmarshal(body, &response)
	if err != nil {
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, err.Error()))
	}
	c.JSON(http.StatusOK, utils.SuccessResponse(response))
}

// @Summary		取消
// @Description	取消
// @Tags			文本转图片
// @Accept			mpfd
// @Produce		json
// @Param			task_id	formData	string	true	"任务ID"
// @Success		200		{string}	string	""
// @Failure		400		{string}	string	""
// @Router			/sdtxt2img/interrupt [post]
func InterruptRequest(c *gin.Context) {

	//url := utils.Cfg.SDWebUI.Url + "/sdapi/v1/interrupt"

	status, _ := redis_func.GetTxt2ImgStatusFromRedis(c.PostForm("task_id"))
	if status == redis_func.SUCCESS {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task is successed"}))
		return
	}
	if status == redis_func.FAILED {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task is failed"}))
		return
	}
	if status == redis_func.PROCESSING {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task is processing,cannot interrupt"}))
		return
	}

	taskId := c.PostForm("task_id")

	err := redis_func.SetTxt2ImgStatusToRedis(taskId, redis_func.CANCEL)
	if err != nil {
		fmt.Println("取消失败", err.Error())
	}

	//data := map[string]string{"task_id": taskId}
	//jsonData, _ := json.Marshal(data)
	//resp, err := services.SubmitPost(url, jsonData)
	//if err != nil {
	//	c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
	//}
	//defer resp.Body.Close()

	c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "interrupt success"}))
}

// @Summary		获取文本转图片结果
// @Description	获取文本转图片结果
// @Tags			文本转图片
// @Accept			mpfd
// @Produce		json
// @Param			task_id	formData	string	true	"任务ID"
// @Success		200		{string}	string	""
// @Failure		400		{string}	string	""
// @Router			/sdtxt2img/getresult [post]
func GetTxt2ImgResult(c *gin.Context) {
	taskId := c.PostForm("task_id")

	status, _ := redis_func.GetTxt2ImgStatusFromRedis(taskId)
	if status == redis_func.PROCESSING {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task is processing"}))
		return
	}
	if status == redis_func.FAILED {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task failed"}))
		return
	}
	if status == redis_func.CANCEL {
		c.JSON(http.StatusOK, utils.SuccessResponse(gin.H{"message": "task is canceled"}))
		return
	}

	result, err := redis_func.GetTxt2ImgTaskFromRedis(taskId)
	if err != nil {
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, "get result failed"))
		return
	}
	var response models.SuccessResponse
	err = json.Unmarshal([]byte(result), &response)
	if err != nil {
		c.JSON(http.StatusInternalServerError, utils.ErrorResponse(http.StatusInternalServerError, "get result failed"))
		return
	}
	for i, img := range response.Images {
		response.Images[i] = utils.ConvertToURL(img)
	}
	c.JSON(http.StatusOK, utils.SuccessResponse(response))
}

// @Summary		批量获取文本转图片状态
// @Description	批量获取文本转图片状态
// @Tags			文本转图片
// @Accept			mpfd
// @Produce		json
// @Param			task_ids	formData	string	true	"任务ID数组"
// @Success		200			{string}	string	""
// @Failure		400			{string}	string	""
// @Router			/sdtxt2img/gettxt2imgstatus [post]
func GetTxt2ImgStatus(c *gin.Context) {
	taskIdsStr := c.PostForm("task_ids")
	taskIds := strings.Split(taskIdsStr, ",")
	var IDStatus []map[string]string
	for _, taskId := range taskIds {
		statusInt, _ := redis_func.GetTxt2ImgStatusFromRedis(taskId)
		var status string
		switch statusInt {
		case redis_func.PROCESSING:
			status = "processing" // 正在生成
		case redis_func.FAILED:
			status = "failed" // 生成失败
		case redis_func.SUCCESS:
			status = "success" // 生成成功
		case redis_func.CANCEL:
			status = "cancel" // 取消
		case redis_func.QUEUED:
			status = "queued" // 排队中
		default:
			status = "unknown" // 未知
		}
		IDStatus = append(IDStatus, map[string]string{"task_id": taskId, "status": status})
	}
	c.JSON(http.StatusOK, utils.SuccessResponse(IDStatus))

}
